#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: my_videowriter.py
@time: 2020/5/25 12:04
@desc:
'''

import cv2

class MyVideoWriter:
    def __init__(self,save_path,height=None,width=None):
        self.save_path = save_path
        self.height = height
        self.width = width
        self.fourcc = cv2.VideoWriter_fourcc(*'MP4V')
        self.cur_frame_id = -1
        self.videoWriter = None

    def save(self,frame):
        self.cur_frame_id += 1
        if self.cur_frame_id == 0:
            if self.height is None or self.width is None:
                self.height,self.width,_ = frame.shape
            self.videoWriter = cv2.VideoWriter(self.save_path,self.fourcc, 20, (self.width, self.height))
        self.videoWriter.write(frame)

    def release(self):
        if self.videoWriter is not None:
            self.videoWriter.release()



